package com.study.cloud.apilog.response;

import cn.hutool.core.bean.BeanUtil;
import com.aliyun.openservices.aliyun.log.producer.Producer;
import com.aliyun.openservices.aliyun.log.producer.errors.LogSizeTooLargeException;
import com.aliyun.openservices.aliyun.log.producer.errors.ProducerException;
import com.aliyun.openservices.aliyun.log.producer.errors.TimeoutException;
import com.aliyun.openservices.log.common.LogItem;
import com.study.cloud.apilog.GatewayApiLog;
import com.study.cloud.apilog.MyCachedBodyOutputMessage;
import com.study.cloud.apilog.entity.SLSProperties;
import com.study.cloud.apilog.sls.SampleCallback;
import com.study.cloud.utils.ProducerUtil;
import com.study.cloud.utils.TraceIdUtil;
import org.apache.commons.lang3.StringUtils;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @Date 2022/5/17
 * @Author Cliff
 * desc
 */
public abstract class AbstractResponseBody implements ResponseBodyService, DisposableBean, ApplicationContextAware {

    private Logger logger = LoggerFactory.getLogger(AbstractResponseBody.class);
    @Autowired
    SLSProperties slsProperties;
    private Producer producer;
    private ExecutorService executorService = ProducerUtil.getExecutorService();;
    final AtomicLong completed = new AtomicLong(0);
    private ApplicationContext applicationContext;
    private String profileActive = "";


    /**
     * 请求装饰器，重新计算 headers
     *
     * @param exchange
     * @param headers
     * @param outputMessage
     * @return
     */
    public ServerHttpRequestDecorator requestDecorate(ServerWebExchange exchange, HttpHeaders headers,
                                                      MyCachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public HttpHeaders getHeaders() {
                long contentLength = headers.getContentLength();
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll(super.getHeaders());
                if (contentLength > 0) {
                    httpHeaders.setContentLength(contentLength);
                } else {
                    httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                }
                return httpHeaders;
            }

            @Override
            public Flux<DataBuffer> getBody() {
                return outputMessage.getBody();
            }
        };
    }

    /**
     * 记录响应日志
     * 通过 DataBufferFactory 解决响应体分段传输问题。
     */
    public ServerHttpResponseDecorator recordResponseLog(ServerWebExchange exchange, GatewayApiLog gatewayLog) {
        ServerHttpResponse response = exchange.getResponse();
        DataBufferFactory bufferFactory = response.bufferFactory();
        return new ServerHttpResponseDecorator(response) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {

                if (body instanceof Flux) {
                    LocalDateTime responseTime = LocalDateTime.now();
                    gatewayLog.setTimeEnd(responseTime);
                    // 计算执行时间
                    long executeTime = responseTime.atZone(ZoneId.systemDefault()).toInstant().toEpochMilli() -
                            gatewayLog.getTimeStart().atZone(ZoneId.systemDefault()).toInstant().toEpochMilli();

                    gatewayLog.setTimeExecute(String.valueOf(executeTime));
                    gatewayLog.setStatus(String.valueOf(response.getStatusCode().value()));

                    // 获取响应类型，如果是 json 就打印
                    String originalResponseContentType = exchange.getAttribute(ServerWebExchangeUtils.ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR);

                    if (StringUtils.isNotBlank(originalResponseContentType)
                            && originalResponseContentType.contains(MediaType.APPLICATION_JSON_VALUE)) {

                        Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                        return super.writeWith(fluxBody.buffer().map(dataBuffers -> {

                            // 合并多个流集合，解决返回体分段传输
                            DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                            DataBuffer join = dataBufferFactory.join(dataBuffers);
                            byte[] content = new byte[join.readableByteCount()];
                            join.read(content);

                            // 释放掉内存
                            DataBufferUtils.release(join);
                            String responseResult = new String(content, StandardCharsets.UTF_8);
                            gatewayLog.setResponseData(responseResult);

                            return bufferFactory.wrap(content);
                        }));
                    }
                }
                return super.writeWith(body);
            }
        };
    }


    public Mono<Void> getThen(ServerWebExchange exchange, GatewayFilterChain chain, GatewayApiLog gatewayLog, ServerHttpRequest mutableReq, ServerHttpResponseDecorator decoratedResponse) {
        return chain.filter(exchange.mutate().request(mutableReq).response(decoratedResponse).build())
                .then(Mono.fromRunnable(() -> {
                    dealRequestPath(exchange, gatewayLog);
                    // 打印日志
                    writeAccessLog(gatewayLog);
                    // 清除TraceID
                    TraceIdUtil.clear();
                }));
    }


    /**
     * 设置请求路径
     *
     * @param exchange
     * @param gatewayLog
     */
    private void dealRequestPath(ServerWebExchange exchange, GatewayApiLog gatewayLog) {
        URI targetUri = getGatewayURI(exchange);
        gatewayLog.setRequestPath(targetUri.getPath());
        gatewayLog.setRequestPathFull(targetUri.toString());
    }

    /**
     * 获取正式的请求url
     *
     * @param exchange
     * @return
     */
    private URI getGatewayURI(ServerWebExchange exchange) {
        return exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
    }

    /**
     * @param gatewayApiLog 网关日志
     */
    private void writeAccessLog(GatewayApiLog gatewayApiLog) {

        producer = ProducerUtil.getInstance(slsProperties);

        LogItem logItem = new LogItem();
        BeanUtil.beanToMap(gatewayApiLog,true,false)
                .forEach((key, value) -> {
                    String v;
                    if(value instanceof LocalDateTime){
                        LocalDateTime ldt = (LocalDateTime) value;
                        v= ldt.format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"));
                    }else{
                        v = value == null ? null : (String) value;
                    }

                    logItem.PushBack(key, v);
                });

        // 提交任务
        executorService.submit(() -> {
            try {
                String project = slsProperties.getProjectName();
                String logStore = slsProperties.getLogstoreName();
                producer.send(project, logStore, profileActive,
                        slsProperties.getSource(), logItem,
                        new SampleCallback(project, logStore, logItem, completed));
            } catch (InterruptedException e) {
                logger.warn("当前线程在发送日志期间被中断.");
            } catch (Exception e) {
                // 发送飞书消息 todo
                if (e instanceof LogSizeTooLargeException) {
                    logger.error("日志大小大于最大允许大小, e={}", e);
                } else if (e instanceof TimeoutException) {
                    logger.error("为日志分配内存所花费的时间已超过, e={}", e);
                } else {
                    logger.error("日志发送失败, logItem={}, e=", logItem, e);
                }

            }
        });
        logger.info("发送成功={}", completed.get());
    }


    @Override
    public void destroy() throws Exception {
        // 销毁的时候调用
        executorService.shutdown();

        try {
            producer.close();
        } catch (InterruptedException e) {
            logger.warn("The current thread has been interrupted from close.");
        } catch (ProducerException e) {
            logger.info("Failed to close producer, e=", e);
        }
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
        String[] activeProfiles = applicationContext.getEnvironment().getActiveProfiles();
        if(activeProfiles != null && activeProfiles.length > 0){
            profileActive = activeProfiles[0];
        }

    }
}
